from model import EnsembleMLPCtrlModel, MlpPartialIKModel
from model.data_processing.data_loaders import Batch, NormalizedEndJointData
from model.data_processing.jax_dataloader import NumpyLoader


def train_end_joint():
    data = NormalizedEndJointData("data/json_data/allover.json")
    data_loader = NumpyLoader(
        data,
        batch_size=512,
        shuffle=True,
        num_workers=10,
        pin_memory=False,
    )
    model = MlpPartialIKModel(in_dim=3, out_dim=6, dropout_rate=0.1, weight_decay=1e-4)
    model.train(data_loader, 20000, save_as="mlp_allover_end_joint.ckp", val_data=None)


def train_ik_6d():

    class NormalizedEJ6D(NormalizedEndJointData):
        """
        Normaized (x, y, z, rx, ry, rz) -> (j1:j6)
        """

        def __init__(self, **kwargs):
            super().__init__(**kwargs)

        def __getitem__(self, idx):
            return Batch(self.ends[idx], self.joints[idx])

    data = NormalizedEJ6D(filename="data/json_data/allover.json")
    data_loader = NumpyLoader(
        data,
        batch_size=512,
        shuffle=True,
        num_workers=10,
        pin_memory=False,
    )
    model = MlpPartialIKModel(in_dim=6, out_dim=6, dropout_rate=0.1, weight_decay=1e-4)
    model.train(data_loader, 20000, save_as="mlp_6d_end_joint.ckp", val_data=None)


def train_ensemble_end_joint():
    from model.data_processing.data_loaders import NormalizedEndJointData

    data = NormalizedEndJointData("data/json_data/allover.json")
    data_loader = NumpyLoader(
        data,
        batch_size=512,
        shuffle=True,
        num_workers=10,
        pin_memory=False,
    )

    model = EnsembleMLPCtrlModel(in_dim=3, out_dim=6, num_heads=10, weight_decay=1e-4)
    model.train(data_loader, 20000, save_as="ensemble_end_joint.ckp", val_data=None)


if __name__ == "__main__":
    train_end_joint()
    # train_ensemble_end_joint()

    # train_ik_6d()
    # train_orien_pred()

# def train_arm_joint():
#     from data.data_loaders import ArmJointData

#     data = ArmJointData("data/json_data/can0_L20Full.json")
#     data_loader = NumpyLoader(
#         data,
#         batch_size=512,
#         shuffle=True,
#         num_workers=10,
#         pin_memory=False,
#     )

#     model = MLPArmJointModel()
#     model.train(data_loader, 30000, save_as="nn_arm_joint.ckp")


# def train_fourier_end_joint():
#     from data.data_loaders import FourierEndJointData

#     data = FourierEndJointData("data/json_data/can0_L20Full.json")
#     data_loader = NumpyLoader(
#         data,
#         batch_size=512,
#         shuffle=True,
#         num_workers=10,
#         pin_memory=False,
#     )

#     model = MLPFourierEndJointModel()
#     model.train(data_loader, 30000, save_as="nn_end_joint_fourier.ckp")


# def train_normalized_end_joint():
#     from model import MLPNormEndJointModel
#     from data.data_loaders import NormalizedEndJointData

#     data = NormalizedEndJointData("data/json_data/can0_L20Full.json")
#     data_loader = NumpyLoader(
#         data,
#         batch_size=512,
#         shuffle=True,
#         num_workers=10,
#         pin_memory=False,
#     )

#     model = MLPNormEndJointModel()
#     model.train(data_loader, 30000, save_as="nn_end_joint_norm.ckp")

# if __name__ == "__main__":
#     train_end_joint()

#     data = EndJointData("json_data/can0_L20Full.json")
#     data_loader = NumpyLoader(
#             data,
#             batch_size=512,
#             shuffle=True,
#             num_workers=10,
#             pin_memory= False,
#         )

#     self = GripJointModel()
#     # self.load_ckpt("data/weights/nn_3d_6d_LN.ckp")
#     self.train(data_loader, 50000, save_as="nn_3d_6d_129535.ckp")
